#include "hybrid_client.h"

int check_pubkey_file()
{
// 检查公钥文件是否存在
#ifdef _WIN32
  if ((0 == _access("./KEY/kyberPubKey.bin", F_OK)) && (0 == access("./KEY/SMPubKey.bin", F_OK)))
    return 1;
  else
    return 0;
#else
  if ((0 == access("./KEY/kyberPubKey.bin", F_OK)) && (0 == access("./KEY/SMPubKey.bin", F_OK)))
    return 1;
  else
    return 0;
#endif
}

int hybrid_encrypt(unsigned char *message, unsigned char *cipherText, int *cipherLen)
{
  if (0 == check_pubkey_file())
    return -1;

  unsigned char msgKyber[ENCRYPT_KEY_LEN]; // Kyber要处理的原始内容
  unsigned char msgSM[ENCRYPT_KEY_LEN];    // SM2要处理的原始内容
  FILE *inKyberPubFile;                    // Kyber读取公钥的文件指针
  FILE *inSMPubFile;                       // SM2读取公钥的文件指针

  // 函数时间
  struct timeval li_enc_start, li_enc_end;
  double enc_time;

  // Kyber
  unsigned char kyberPubKey[KYBER_INDCPA_PUBLICKEYBYTES]; // Kyber的公钥
  unsigned char coins[KYBER_SYMBYTES];                    // Kyber的随机数

  memcpy(msgKyber, message, ENCRYPT_KEY_LEN);
  // printf("原始密钥Kyber:");
  // for (int i = 0; i < ENCRYPT_KEY_LEN; i++)
  // {
  //   printf("%02X", msgKyber[i]);
  // }

  // read public key
  if (!(inKyberPubFile = fopen("./KEY/kyberPubKey.bin", "rb")))
    return -1;
  fread(kyberPubKey, sizeof(char), KYBER_INDCPA_PUBLICKEYBYTES, inKyberPubFile);
  fclose(inKyberPubFile);

  kyber_randombytes(coins, KYBER_SYMBYTES);

  // 获取开始时间
  gettimeofday(&li_enc_start,NULL);

  indcpa_enc(cipherText, msgKyber, kyberPubKey, coins);

  // 获取结束时间并计算函数执行时间
  gettimeofday(&li_enc_end,NULL);
	enc_time = 1000000*(li_enc_end.tv_sec-li_enc_start.tv_sec)+(li_enc_end.tv_usec-li_enc_start.tv_usec);
	printf("\nKyber加密时间: %.0fus\n", enc_time);
  // End Kyber

  // SM2
  SM2_KEY SMPubKey; // SM2的公钥
  size_t SMLen;     // SM2加密后密文长度

  if (!(inSMPubFile = fopen("./KEY/SMPubKey.bin", "rb")))
    return -1;
  fread(&SMPubKey.public_key, sizeof(SM2_POINT), 1, inSMPubFile);
  fclose(inSMPubFile);

  memcpy(msgSM, message + ENCRYPT_KEY_LEN, ENCRYPT_KEY_LEN);

  // printf("\n原始密钥SM2:");
  // for (int i = 0; i < ENCRYPT_KEY_LEN; i++)
  // {
  //   printf("%02X", msgSM[i]);
  // }
  // printf("\n");

  // 获取开始时间
  gettimeofday(&li_enc_start,NULL);

  sm2_encrypt(&SMPubKey, msgSM, ENCRYPT_KEY_LEN, cipherText + KYBER_INDCPA_BYTES, &SMLen);

  // 获取结束时间并计算函数执行时间
  gettimeofday(&li_enc_end,NULL);
	enc_time = 1000000*(li_enc_end.tv_sec-li_enc_start.tv_sec)+(li_enc_end.tv_usec-li_enc_start.tv_usec);
	printf("\nSM2加密时间: %.0fus\n", enc_time);
  // End SM2

  // printf("SM2 Length:%d\n", SMLen);

  *cipherLen = KYBER_INDCPA_BYTES + SMLen; // 传回密文总长度

  return 1;
}

int cal_sm3(unsigned char *key, unsigned char *sm3hash)
{
  SM3_CTX sm3_ctx; // SM3结构

  sm3_init(&sm3_ctx);
  sm3_update(&sm3_ctx, key, ORIGIN_KEY_LEN);
  sm3_finish(&sm3_ctx, sm3hash);
}

int main()
{
  unsigned char originKey[ORIGIN_KEY_LEN] = {0};                      // 原始密钥，本地生成，长度64
  unsigned char originHash[HASH_LEN];                                 // 原始密钥的哈希，本地计算，长度32
  unsigned char cipherKey[KYBER_INDCPA_BYTES + SM2_CIPHER_LEN] = {0}; // 密文，本地计算
  int cipherLen = 0;                                                  // 密文长度，本地计算

  // 函数时间
  struct timeval li_func_start, li_func_end;
  double func_time;

  kyber_randombytes(originKey, ORIGIN_KEY_LEN);

  printf("原始密钥(Kyber||SM2):\n");
  for (int i = 0; i < ORIGIN_KEY_LEN; i++)
  {
    printf("%02X", originKey[i]);
    if (i == 32)
      printf(" ");
  }
  printf("\n");

  // 获取开始时间
  gettimeofday(&li_func_start,NULL);

  cal_sm3(originKey, originHash);

  // 获取结束时间并计算函数执行时间
  gettimeofday(&li_func_end,NULL);
	func_time = 1000000*(li_func_end.tv_sec-li_func_start.tv_sec)+(li_func_end.tv_usec-li_func_start.tv_usec);
	printf("\n计算SM3时间: %.0fus\n", func_time);

  printf("原始密钥Hash:\n");
  for (int i = 0; i < HASH_LEN; i++)
  {
    printf("%02X", originHash[i]);
  }
  printf("\n");

  hybrid_encrypt(originKey, cipherKey, &cipherLen);

  printf("加密后密钥(部分):");
  for (int i = 0; i < ORIGIN_KEY_LEN; i++)
  {
    printf("%02X", cipherKey[i]);
  }
  printf("\n");

  printf("加密后密钥长度:%d\n", cipherLen);

  int sockfd, n;
  char sendline[MAXLINE] = {0};
  struct sockaddr_in servaddr;

  if ((sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0)
  {
    printf("创建Socket套接字失败: %s(errno: %d)\n", strerror(errno), errno);
    exit(0);
  }

  servaddr.sin_addr.s_addr = inet_addr("127.0.0.1");
  servaddr.sin_family = AF_INET;
  servaddr.sin_port = htons(7658);

  if (connect(sockfd, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0)
  {
    printf("连接服务器失败: %s(errno: %d)\n", strerror(errno), errno);
    exit(0);
  }

  printf("发送加密后数据到对端节点\n");

  // cipherLen||originHash||cipherKey
  // 密文长度||原始密钥哈希值||密文
  memcpy(sendline, &cipherLen, sizeof(int));
  memcpy(sendline + sizeof(int), originHash, HASH_LEN);
  memcpy(sendline + sizeof(int) + HASH_LEN, cipherKey, cipherLen);

  if (send(sockfd, sendline, sizeof(int) + HASH_LEN + cipherLen, 0) < 0)
  {
    printf("send msg error: %s(errno: %d)\n", strerror(errno), errno);
    exit(0);
  }

  close(sockfd);
  exit(0);

  return 1;
}